import torch

from comfy.patcher_extension import WrappersMP
from comfy.model_patcher import ModelPatcher

from .tkg import apply_tkg_noise, COLOR_SETS, COLOR_SET_MAP


def patch_first_sample(
    shift: float,
    grid_factor: int,
    channels: list[float],
):
    def sampler_sample_wrapper(
        sampler,
        model_wrap: ModelPatcher,
        sigmas: torch.Tensor,
        extra_args: dict,
        callback,
        noise: torch.Tensor,
        latent_image=None,
        denoise_mask=None,
        disable_pbar=False,
    ):
        assert len(channels) == noise.shape[1], (
            f"Number of channels ({len(channels)}) must match latent channels ({noise.shape[1]})."
        )

        new_noise = apply_tkg_noise(
            noise,
            shift=shift,
            delta_shift=0.1,
            std_dev=0.5,
            factor=grid_factor,
            channels=channels,
        ).to(noise.dtype)

        return sampler(
            model_wrap,
            sigmas,
            extra_args,
            callback,
            new_noise,
            latent_image,
            denoise_mask,
            disable_pbar,
        )

    return sampler_sample_wrapper


class ApplyTKGChromaKeySDXL:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": (
                    "MODEL",
                    {},
                ),
                "color": (
                    [c.name for c in COLOR_SETS],
                    {
                        "default": "green",
                    },
                ),
                "shift": (
                    "FLOAT",
                    {
                        "default": 0.11,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                    },
                ),
                "grid_factor": (
                    "INT",
                    {
                        "default": 8,
                        "min": 1,
                        "max": 16,
                        "step": 1,
                    },
                ),
            }
        }

    RETURN_TYPES = ("MODEL",)
    OUTPUT_TOOLTIPS = ("Output latents with TKG noise applied",)
    FUNCTION = "apply_tkg"

    CATEGORY = "latent"
    DESCRIPTION = "Apply TKG chroma key noise to latents using channel-wise mean shift from predefined color sets."

    CATEGORY = "latent"

    def apply_tkg(
        self,
        model: ModelPatcher,
        color: str,
        shift: float = 0.11,
        grid_factor: int = 8,
    ):
        color_set = COLOR_SET_MAP[color]

        _model = model.clone()

        _model.add_wrapper(
            WrappersMP.SAMPLER_SAMPLE,
            patch_first_sample(
                shift=shift,
                grid_factor=grid_factor,
                channels=color_set.channels,
            ),
        )

        return (_model,)


class ApplyTKGChromaKeyAdvanced:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": (
                    "MODEL",
                    {},
                ),
                "channels": (
                    "STRING",
                    {
                        "default": "0, 1, 1, 0",
                    },
                ),
                "shift": (
                    "FLOAT",
                    {
                        "default": 0.11,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                    },
                ),
                "grid_factor": (
                    "INT",
                    {
                        "default": 8,
                        "min": 1,
                        "max": 16,
                        "step": 1,
                    },
                ),
            }
        }

    RETURN_TYPES = ("MODEL",)
    OUTPUT_TOOLTIPS = ("Output latents with TKG noise applied",)
    FUNCTION = "apply_tkg"

    CATEGORY = "latent"
    DESCRIPTION = "Apply TKG chroma key noise to latents using a channel-wise mean shift with manually specified channels."

    CATEGORY = "latent"

    def apply_tkg(
        self,
        model: ModelPatcher,
        channels: str,
        shift: float = 0.11,
        grid_factor: int = 8,
    ):
        channels_list = self.validate_channels(channels)

        _model = model.clone()

        _model.add_wrapper(
            WrappersMP.SAMPLER_SAMPLE,
            patch_first_sample(
                shift=shift,
                grid_factor=grid_factor,
                channels=channels_list,
            ),
        )

        return (_model,)

    def validate_channels(self, channels: str) -> list[float]:
        try:
            channel_list = [float(c.strip()) for c in channels.split(",")]
            return channel_list
        except ValueError:
            raise ValueError(
                "Channels must be a comma-separated list of four numbers. e.g. '0, 1, 1, 0'"
            )
